#pragma once
#include "stdafx.h"
#include <vector>
#include <string>

namespace TcpProxy
{

    struct TCP_PACKET
    {
        TCP_PACKET()
        {
            buffer.len = 0;
            buffer.buf = NULL;
        }
        TCP_PACKET(const char *buf, int len)
        {
            if (len > 0)
            {
                buffer.buf = new char[len];
                buffer.len = len;

                if (buf)
                {
                    memcpy(buffer.buf, buf, len);
                }
            }
            else
            {
                buffer.buf = NULL;
                buffer.len = 0;
            }
        }

        WSABUF &operator()()
        {
            return buffer;
        }

        void free()
        {
            if (buffer.buf)
            {
                delete[] buffer.buf;
            }
        }

        WSABUF buffer;
    };

    typedef std::vector<TCP_PACKET> tPacketList;

    enum OV_TYPE
    {
        OVT_ACCEPT,
        OVT_CONNECT,
        OVT_CLOSE,
        OVT_SEND,
        OVT_RECEIVE
    };

    struct OV_DATA
    {
        OV_DATA()
        {
            memset(&ol, 0, sizeof(ol));
        }
        ~OV_DATA()
        {
            for (tPacketList::iterator it = packetList.begin(); it != packetList.end(); it++)
            {
                it->free();
            }
        }

        OVERLAPPED ol;
        LIST_ENTRY entry;
        LIST_ENTRY entryEventList;
        NFAPI_NS ENDPOINT_ID id;
        OV_TYPE type;
        tPacketList packetList;

        SOCKET socket;
        DWORD dwTransferred;
        int error;
    };

    enum PROXY_STATE
    {
        PS_NONE,
        PS_AUTH,
        PS_AUTH_NEGOTIATION,
        PS_CONNECT,
        PS_CONNECTED,
        PS_ERROR,
        PS_CLOSED
    };

    enum PROXY_TYPE
    {
        PROXY_NONE,
        PROXY_SOCKS5
    };

    struct SOCKET_DATA
    {
        SOCKET_DATA()
        {
            socket = INVALID_SOCKET;
            disconnected = false;
            receiveInProgress = false;
            sendInProgress = false;
            disconnect = false;
            closed = false;
        }
        ~SOCKET_DATA()
        {
            if (socket != INVALID_SOCKET)
            {
                closesocket(socket);
            }
            for (tPacketList::iterator it = packetList.begin(); it != packetList.end(); it++)
            {
                it->free();
            }
        }

        SOCKET socket;
        bool disconnected;
        bool receiveInProgress;
        bool sendInProgress;
        bool disconnect;
        bool closed;
        tPacketList packetList;
    };

    struct PROXY_DATA
    {
        PROXY_DATA()
        {
            id = 0;
            proxyState = PS_NONE;
            suspended = false;
            offline = false;
            memset(&connInfo, 0, sizeof(connInfo));
            refCount = 1;
            proxyType = PROXY_NONE;
            proxyAddressLen = 0;
        }
        ~PROXY_DATA()
        {
        }

        NFAPI_NS ENDPOINT_ID id;

        SOCKET_DATA inSocket;
        SOCKET_DATA outSocket;

        PROXY_STATE proxyState;
        NFAPI_NS NF_TCP_CONN_INFO connInfo;

        PROXY_TYPE proxyType;
        char proxyAddress[NF_MAX_ADDRESS_LENGTH];
        int proxyAddressLen;

        std::string userName;
        std::string userPassword;

        bool suspended;
        bool offline;

        int refCount;
        AutoCriticalSection lock;

    private:
        PROXY_DATA &operator=(const PROXY_DATA &v)
        {
            return *this;
        }
    };
};
namespace UdpProxy
{
    typedef std::vector<char> tBuffer;
    enum OV_TYPE
    {
        OVT_CONNECT,
        OVT_TCP_SEND,
        OVT_TCP_RECEIVE,
        OVT_UDP_SEND,
        OVT_UDP_RECEIVE
    };

    struct OV_DATA
    {
        OV_DATA()
        {
            memset(&ol, 0, sizeof(ol));
        }

        OVERLAPPED ol;
        ENDPOINT_ID id;
        OV_TYPE type;
        char remoteAddress[NF_MAX_ADDRESS_LENGTH];
        int remoteAddressLen;
        tBuffer buffer;
    };

    typedef std::set<OV_DATA *> tOvDataSet;

    enum PROXY_STATE
    {
        PS_NONE,
        PS_AUTH,
        PS_AUTH_NEGOTIATION,
        PS_UDP_ASSOC,
        PS_CONNECTED,
        PS_CLOSED
    };

    struct UDP_PACKET
    {
        char remoteAddress[NF_MAX_ADDRESS_LENGTH];
        int remoteAddressLen;
        tBuffer buffer;
    };

    typedef std::list<UDP_PACKET *> tPacketList;

    struct PROXY_DATA
    {
        PROXY_DATA()
        {
            tcpSocket = INVALID_SOCKET;
            udpSocket = INVALID_SOCKET;
            proxyState = PS_NONE;
            memset(remoteAddress, 0, sizeof(remoteAddress));
            remoteAddressLen = 0;
            udpRecvStarted = false;
            ipFamily = AF_INET;
        }
        ~PROXY_DATA()
        {
            if (tcpSocket != INVALID_SOCKET)
            {
                shutdown(tcpSocket, SD_BOTH);
                closesocket(tcpSocket);
            }
            if (udpSocket != INVALID_SOCKET)
            {
                closesocket(udpSocket);
            }
            while (!udpSendPackets.empty())
            {
                tPacketList::iterator it = udpSendPackets.begin();
                delete (*it);
                udpSendPackets.erase(it);
            }
        }

        SOCKET tcpSocket;
        SOCKET udpSocket;

        PROXY_STATE proxyState;

        char remoteAddress[NF_MAX_ADDRESS_LENGTH];
        int remoteAddressLen;

        std::string userName;
        std::string userPassword;

        tPacketList udpSendPackets;

        bool udpRecvStarted;

        u_short ipFamily;
    };
    typedef std::map<ENDPOINT_ID, PROXY_DATA *> tSocketMap;
};